package com.yys.apm.boot;

import javassist.*;

/**
 * agent 装载器
 *
 * @author Yingsheng Ye
 */
public class AgentBootLoader {

    private final String className;

    private final ClassLoader loader;

    private final CtClass ctClass;

    private final byte[] classFileBuffer;

    //代理方法标示
    private final static String AGENT = "$$agent";


    public AgentBootLoader(ClassLoader loader,
                    String className,
                    byte[] classFileBuffer,
                    CtClass ctClass
    ) {
        this.className = className;
        this.loader = loader;
        this.ctClass = ctClass;
        this.classFileBuffer = classFileBuffer;

    }

    /**
     * 代理方法
     *
     * @param ctMethod    CtMethod
     * @param sourceBuild MethodSourceBuild
     * @throws CannotCompileException 异常
     */
    public void agentMethod(CtMethod ctMethod, MethodSourceBuild sourceBuild) throws CannotCompileException {
        String methodName = ctMethod.getName();
        // 重构被代理的方法名称

        // 基于原方法复制生成代理方法
        CtMethod agentMethod = CtNewMethod.copy(ctMethod, methodName, ctClass, null);
        agentMethod.setName(methodName + AGENT);
        ctClass.addMethod(agentMethod);

        // 原方法重置为代理执行
        ctMethod.setBody(sourceBuild.buildSource(ctMethod));
    }

    public String getClassName() {
        return className;
    }

    public ClassLoader getLoader() {
        return loader;
    }

    public CtClass getCtClass() {
        return ctClass;
    }

    public byte[] getClassFileBuffer() {
        return classFileBuffer;
    }

    public static class MethodSourceBuild {

        private final static String SOURCE_TPL = "{\n"
                + "%s"
                + "        Object result=null;\n"
                + "       try {\n"
                + "            result=($w)%s" + AGENT + "($$);\n"
                + "        } catch (Throwable e) {\n"
                + "%s"
                + "            throw e;\n"
                + "        }finally{\n"
                + "%s"
                + "        }\n"
                + "        return ($r) result;\n"
                + "}\n";

        private final static String VOID_SOURCE_TPL = "{\n"
                + "%s"
                + "        try {\n"
                + "            %s" + AGENT + "($$);\n"
                + "        } catch (Throwable e) {\n"
                + "%s"
                + "            throw e;\n"
                + "        }finally{\n"
                + "%s"
                + "        }\n"
                + "}\n";

        //开始源码块
        private String startSrc;

        //结束源码块
        private String endSrc;

        //异常源码块
        private String errorSrc;

        String buildSource(CtMethod method) {
            try {
                String template = method.getReturnType().getName().equals("void") ? VOID_SOURCE_TPL : SOURCE_TPL;
                return String.format(template,
                        startSrc == null ? "" : startSrc,
                        method.getName(),
                        errorSrc == null ? "" : errorSrc,
                        endSrc == null ? "" : endSrc);
            } catch (NotFoundException e) {
                throw new RuntimeException(e);
            }
        }

        public String getStartSrc() {
            return startSrc;
        }

        public void setStartSrc(String startSrc) {
            this.startSrc = startSrc;
        }

        public String getEndSrc() {
            return endSrc;
        }

        public void setEndSrc(String endSrc) {
            this.endSrc = endSrc;
        }

        public String getErrorSrc() {
            return errorSrc;
        }

        public void setErrorSrc(String errorSrc) {
            this.errorSrc = errorSrc;
        }

    }

}
